package com.hmtech.ayp.common.s3.service;

import cn.hutool.core.io.IoUtil;
import com.amazonaws.HttpMethod;
import com.amazonaws.auth.AWSCredentials;
import com.amazonaws.auth.AWSStaticCredentialsProvider;
import com.amazonaws.auth.BasicAWSCredentials;
import com.amazonaws.services.s3.AmazonS3;
import com.amazonaws.services.s3.AmazonS3ClientBuilder;
import com.amazonaws.services.s3.model.*;
import com.amazonaws.services.s3.transfer.TransferManager;
import com.amazonaws.services.s3.transfer.Upload;
import com.amazonaws.services.securitytoken.AWSSecurityTokenService;
import com.amazonaws.services.securitytoken.AWSSecurityTokenServiceClientBuilder;
import com.amazonaws.services.securitytoken.model.AssumeRoleRequest;
import com.amazonaws.services.securitytoken.model.AssumeRoleResult;
import com.amazonaws.services.securitytoken.model.Credentials;
import com.hmtech.ayp.common.s3.S3Utils;
import com.hmtech.ayp.common.s3.domain.AwsTempCredentials;
import com.hmtech.ayp.common.s3.domain.S3UploadResult;
import com.hmtech.ayp.common.s3.enums.S3ErrorEnum;
import com.hmtech.ayp.common.s3.exception.S3Exception;
import com.hmtech.ayp.common.s3.properties.AwsS3Properties;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import java.io.*;
import java.net.MalformedURLException;
import java.net.URL;
import java.util.ArrayList;
import java.util.Date;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.TimeUnit;

/**
 * @author 闫家洪（yanjiahong@vcredit.com）
 * @since 2023/12/19
 */
@Component
@RequiredArgsConstructor(onConstructor = @__(@Autowired))
@Slf4j
public class IS3ServiceImpl implements IS3Service {
    private final AwsS3Properties s3Properties;

    private AmazonS3 s3client;
    private AWSSecurityTokenService stsClient;

    @PostConstruct
    public void init() {
        AWSCredentials credentials = new BasicAWSCredentials(s3Properties.getAccessKeyId(),
                s3Properties.getAccessSecretKey());
        s3client = AmazonS3ClientBuilder.standard()
                .withCredentials(new AWSStaticCredentialsProvider(credentials))
                .withRegion(s3Properties.getRegion()).build();

        // 创建STS客户端
        stsClient = AWSSecurityTokenServiceClientBuilder.standard()
                .withCredentials(new AWSStaticCredentialsProvider(credentials))
                .withRegion(s3Properties.getRegion())
                .build();
    }

    @Override
    public S3UploadResult upload(String filename, byte[] data) throws S3Exception {
        String uniqueFilename = S3Utils.generateUniqueFilename(filename);
        String key = S3Utils.getS3KeyFromUniqueFilename(uniqueFilename);
        ByteArrayInputStream bais = new ByteArrayInputStream(data);
        PutObjectRequest request = new PutObjectRequest(s3Properties.getBucketName(), key, bais, null);
        return getS3UploadResult(key, request);
    }

    @Override
    public S3UploadResult upload(String localFilepath) throws S3Exception {
        String uniqueFilename = S3Utils.generateUniqueFilename(localFilepath);
        String key = S3Utils.getS3KeyFromUniqueFilename(uniqueFilename);
        File file = new File(localFilepath);
        PutObjectRequest request = new PutObjectRequest(s3Properties.getBucketName(), key, file);
        return getS3UploadResult(key, request);
    }

    @Override
    public void download(String s3path, String dstDirPath) throws S3Exception {
        PresignedUrlDownloadResult result = doDownload(s3path);
        InputStream is = result.getS3Object().getObjectContent();
        String key = S3Utils.getKeyFromS3path(s3path);
        String dstFilePath = new File(new File(dstDirPath), key).getAbsolutePath();
        FileOutputStream fos = null;
        try {
            fos = new FileOutputStream(dstFilePath);
            IoUtil.copy(is, fos);
        } catch (FileNotFoundException e) {
            throw new RuntimeException(e);
        } finally {
            if (null != fos) {
                try {
                    fos.close();
                } catch (IOException e) {
                    log.error("文件关闭失败！", e);
                }
            }
        }
    }

    @Override
    public byte[] download(String s3path) throws S3Exception {
        PresignedUrlDownloadResult result = doDownload(s3path);
        S3ObjectInputStream s3Ins = result.getS3Object().getObjectContent();
        return IoUtil.readBytes(s3Ins);
    }

    private PresignedUrlDownloadResult doDownload(String s3path) {
        String presignedUrl = genPreSignedUrlForGet(s3path);
        URL url;
        try {
            url = new URL(presignedUrl);
        } catch (MalformedURLException e) {
            throw new S3Exception(S3ErrorEnum.CREATE_URL_FAIL, e);
        }
        PresignedUrlDownloadRequest downloadRequest = new PresignedUrlDownloadRequest(url);
        return s3client.download(downloadRequest);
    }

    private S3UploadResult getS3UploadResult(String key, PutObjectRequest request) {
        s3client.putObject(request);
        S3UploadResult result = new S3UploadResult();
        result.setKey(key);
        String s3path = S3Utils.getS3Path(s3Properties.getBucketName(), key);
        result.setS3path(s3path);
        result.setPreSignedGetUrl(genPreSignedUrlForGet(s3path));
        return result;
    }

    @Override
    public S3UploadResult syncUploadConcurrently(String localFilepath) throws S3Exception {
        TransferManager transferManager = new TransferManager(s3client);
        try {
            File file = new File(localFilepath);
            String uniqueFilename = S3Utils.generateUniqueFilename(localFilepath);
            String key = S3Utils.getS3KeyFromUniqueFilename(uniqueFilename);
            Upload upload = transferManager.upload(s3Properties.getBucketName(), key, file);
            upload.waitForCompletion();

            S3UploadResult result = new S3UploadResult();
            result.setOriginalFilepath(localFilepath);
            result.setOriginalFilename(file.getName());
            result.setKey(key);
            String s3path = S3Utils.getS3Path(s3Properties.getBucketName(), key);
            result.setPreSignedGetUrl(genPreSignedUrlForGet(s3path));
            result.setS3path(s3path);
        } catch (Exception e) {
            throw new S3Exception(S3ErrorEnum.S3_UPLOAD_ERROR, e);
        } finally {
            transferManager.shutdownNow();
        }
        return null;
    }

    @Override
    public List<S3UploadResult> syncUploadConcurrently(List<String> localFilepathList) throws S3Exception {
        TransferManager transferManager = new TransferManager(s3client);
        List<S3UploadResult> results = new ArrayList<>();
        List<Upload> uploads = new ArrayList<>();
        try {
            for (String filepath : localFilepathList) {
                File file = new File(filepath);
                String uniqueFilename = S3Utils.generateUniqueFilename(filepath);
                String key = S3Utils.getS3KeyFromUniqueFilename(uniqueFilename);
                Upload upload = transferManager.upload(s3Properties.getBucketName(), key, file);
                uploads.add(upload);

                S3UploadResult result = new S3UploadResult();
                result.setOriginalFilepath(filepath);
                result.setOriginalFilename(file.getName());
                result.setKey(key);
                String s3path = S3Utils.getS3Path(s3Properties.getBucketName(), key);
                result.setS3path(s3path);
                result.setPreSignedGetUrl(genPreSignedUrlForGet(s3path));
                results.add(result);
            }
            for (Upload upload : uploads) {
                upload.waitForUploadResult();
            }
        } catch (Exception e) {
            throw new S3Exception(S3ErrorEnum.S3_UPLOAD_ERROR, e);
        } finally {
            transferManager.shutdownNow();
        }
        return results;
    }

    @Override
    public String genPreSignedUrl(String s3path, HttpMethod method, long expiresInSeconds) throws S3Exception {
        // 提取S3 Key（不含bucketName）
        String key = S3Utils.getKeyFromS3path(s3path);
        // 设置过期时间
        Date expiration = new Date(System.currentTimeMillis() + expiresInSeconds);
        GeneratePresignedUrlRequest request = new GeneratePresignedUrlRequest(s3Properties.getBucketName(), key)
                .withMethod(method)
                .withExpiration(expiration);
        URL url = s3client.generatePresignedUrl(request);
        return url.toString();
    }

    @Override
    public String genPreSignedUrl(String s3path, HttpMethod method) throws S3Exception {
        long expireInMillis = TimeUnit.SECONDS.toMillis(s3Properties.getExpireInSeconds());
        return genPreSignedUrl(s3path, method, expireInMillis);
    }

    @Override
    public String genPreSignedUrlForPut(String s3path) throws S3Exception {
        return genPreSignedUrl(s3path, HttpMethod.PUT);
    }

    @Override
    public String genPreSignedUrlForGet(String s3path) throws S3Exception {
        return genPreSignedUrl(s3path, HttpMethod.GET);
    }

    @Override
    public AwsTempCredentials getTempCredentialsForPutObject() throws S3Exception {
        // 创建AssumeRole请求
        String roleSessionName = "session-" + UUID.randomUUID();
        AssumeRoleRequest roleRequest = new AssumeRoleRequest()
                .withRoleArn(s3Properties.getTempS3PutRoleArn())
                .withRoleSessionName(roleSessionName);
        // 执行请求并接收响应
        AssumeRoleResult roleResponse = stsClient.assumeRole(roleRequest);
        Credentials sessionCredentials = roleResponse.getCredentials();
        AwsTempCredentials awsTempCredentials = new AwsTempCredentials();
        awsTempCredentials.setAccessKeyId(sessionCredentials.getAccessKeyId());
        awsTempCredentials.setAccessSecretKey(sessionCredentials.getSecretAccessKey());
        awsTempCredentials.setSessionToken(sessionCredentials.getSessionToken());
        return awsTempCredentials;
    }
}
